import argparse
import cv2
import torch
import os
import numpy as np
from yolov5.utils.general import non_max_suppression, scale_coords
from yolov5.utils.torch_utils import select_device
from yolov5.utils.plots import plot_one_box
from yolov5.models.yolo import Model

class DefectDetector:
    def __init__(self, weights, img_size=640, conf_thres=0.25, iou_thres=0.45, device='', classes=None):
        self.weights = weights
        self.img_size = img_size
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        self.device = select_device(device)
        self.classes = classes
        self.model = self.load_model()
        self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names
        self.colors = [[np.random.randint(0, 255) for _ in range(3)] for _ in self.names]

    def load_model(self):
        """加载YOLOv5模型"""
        model = Model('models/yolov5s.yaml').to(self.device)
        checkpoint = torch.load(self.weights, map_location=self.device)
        model.load_state_dict(checkpoint['model'].float().state_dict())
        model.eval()
        return model

    def preprocess_image(self, img):
        """预处理图像以适应模型输入"""
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (self.img_size, self.img_size))
        img = img.transpose(2, 0, 1)  # HWC to CHW
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(self.device)
        img = img.float() / 255.0  # 归一化到[0,1]
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        return img

    def detect(self, img):
        """执行缺陷检测"""
        original_img = img.copy()
        img = self.preprocess_image(img)

        # 推理
        with torch.no_grad():
            pred = self.model(img)[0]

        # 非极大值抑制
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, classes=self.classes)

        # 处理检测结果
        detections = []
        for i, det in enumerate(pred):
            if len(det):
                # 将坐标缩放到原始图像尺寸
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], original_img.shape).round()

                # 保存检测结果
                for *xyxy, conf, cls in reversed(det):
                    label = f'{self.names[int(cls)]} {conf:.2f}'
                    detections.append({
                        'class': self.names[int(cls)],
                        'confidence': float(conf),
                        'bbox': [int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])]
                    })

                    # 在图像上绘制边界框
                    plot_one_box(xyxy, original_img, label=label, color=self.colors[int(cls)], line_thickness=3)

        return original_img, detections

    def detect_from_image(self, image_path, output_path=None):
        """从图像文件检测缺陷"""
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法读取图像文件: {image_path}")

        result_img, detections = self.detect(img)

        if output_path:
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            cv2.imwrite(output_path, result_img)
            print(f"检测结果已保存至 {output_path}")

        return result_img, detections

    def detect_from_video(self, video_path, output_path=None):
        """从视频文件或摄像头检测缺陷"""
        cap = cv2.VideoCapture(video_path if video_path else 0)
        if not cap.isOpened():
            raise ValueError(f"无法打开视频源: {video_path}")

        if output_path:
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            fps = cap.get(cv2.CAP_PROP_FPS)
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            result_frame, _ = self.detect(frame)

            cv2.imshow('Defect Detection', result_frame)
            if output_path:
                out.write(result_frame)

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

        cap.release()
        if output_path:
            out.release()
        cv2.destroyAllWindows()


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='runs/train/exp/weights/best.pt', help='模型权重路径')
    parser.add_argument('--source', type=str, default='data/images/test', help='图像/视频路径，0表示摄像头')
    parser.add_argument('--img-size', type=int, default=640, help='推理图像尺寸')
    parser.add_argument('--conf-thres', type=float, default=0.4, help='置信度阈值')
    parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU阈值')
    parser.add_argument('--device', default='', help='cuda设备，例如 0 或 0,1,2,3 或 cpu')
    parser.add_argument('--classes', nargs='+', type=int, help='只检测指定类别，例如 0 或 0 2 3')
    parser.add_argument('--output', type=str, default='runs/detect/exp', help='检测结果保存路径')
    return parser.parse_args()


def main(opt):
    detector = DefectDetector(
        weights=opt.weights,
        img_size=opt.img_size,
        conf_thres=opt.conf_thres,
        iou_thres=opt.iou_thres,
        device=opt.device,
        classes=opt.classes
    )

    # 创建输出目录
    os.makedirs(opt.output, exist_ok=True)

    # 判断输入源类型
    if opt.source.isdigit():
        # 摄像头
        detector.detect_from_video(int(opt.source), os.path.join(opt.output, 'output.mp4'))
    elif opt.source.endswith(('.mp4', '.avi', '.mov')):
        # 视频文件
        detector.detect_from_video(opt.source, os.path.join(opt.output, 'output.mp4'))
    else:
        # 图像文件或目录
        if os.path.isdir(opt.source):
            for filename in os.listdir(opt.source):
                if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                    img_path = os.path.join(opt.source, filename)
                    output_path = os.path.join(opt.output, filename)
                    detector.detect_from_image(img_path, output_path)
        else:
            # 单个图像文件
            detector.detect_from_image(opt.source, os.path.join(opt.output, os.path.basename(opt.source)))

if __name__ == '__main__':
    opt = parse_opt()
    main(opt)